import torch
from torch import nn
from typing import List
from varflow.distributions import ConditionalBernoulli
from varflow.flows import FlowLayer
from .mixture import Mixture


class Identity(FlowLayer):
    def log_prob(self, x):
        return self.base_dist.log_prob(x)

    def log_prob_with_sll(self, x):
        return self.base_dist.log_prob_with_sll(x)

    def sample(self, num_samples):
        return self.base_dist.sample(num_samples)

class Skipable(Mixture):
    def __init__(self, flow, sampler):
        super().__init__([Identity(), flow], sampler)
